import os.path as osp
import numpy as np
import torch
import os
import torchvision
import torch.nn.functional as F
from torch.utils.tensorboard import SummaryWriter
from torch_geometric.datasets import TUDataset
from torch_geometric.data import DataLoader
from torch_geometric.nn import GraphConv, TopKPooling,GATConv
from torch_geometric.nn import global_mean_pool as gap, global_max_pool as gmp
import matplotlib.pyplot as plt

class Net(torch.nn.Module):
    def __init__(self,dataset):
        super(Net, self).__init__()

        self.conv1 = GraphConv(dataset.num_features, 128)
        self.pool1 = TopKPooling(128, ratio=0.8)
        self.conv2 = GraphConv(128, 128)
        self.pool2 = TopKPooling(128, ratio=0.8)
        self.conv3 = GraphConv(128, 128)
        self.pool3 = TopKPooling(128, ratio=0.8)

        self.lin1 = torch.nn.Linear(256, 128)
        self.lin2 = torch.nn.Linear(128, 64)
        self.lin3 = torch.nn.Linear(64, dataset.num_classes)

    def forward(self, data):
        x, edge_index, batch = data.x, data.edge_index, data.batch

        x = F.relu(self.conv1(x, edge_index))
        x, edge_index, _, batch, _, _ = self.pool1(x, edge_index, None, batch)
        x1 = torch.cat([gmp(x, batch), gap(x, batch)], dim=1)

        x = F.relu(self.conv2(x, edge_index))
        x, edge_index, _, batch, _, _ = self.pool2(x, edge_index, None, batch)
        x2 = torch.cat([gmp(x, batch), gap(x, batch)], dim=1)

        x = F.relu(self.conv3(x, edge_index))
        x, edge_index, _, batch, _, _ = self.pool3(x, edge_index, None, batch)
        x3 = torch.cat([gmp(x, batch), gap(x, batch)], dim=1)

        x = x1 + x2 + x3

        x = F.relu(self.lin1(x))
        x = F.dropout(x, p=0.5, training=self.training)
        x = F.relu(self.lin2(x))
        x = F.log_softmax(self.lin3(x), dim=-1)

        return x


def train(model,optimizer,train_dataset,train_loader,device):
    model.train()

    loss_all = 0
    for data in train_loader:
        data = data.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, data.y)
        loss.backward()
        loss_all += data.num_graphs * loss.item()
        optimizer.step()
    return loss_all / len(train_dataset)


def tes(model,test_loader,device):
    model.eval()

    correct = 0
    for data in test_loader:
        data = data.to(device)
        pred = model(data).max(dim=1)[1]
        correct += pred.eq(data.y).sum().item()
    return correct / len(test_loader.dataset)

def PoolingWithGCN(DataSetName,times,EPOCHES):
    writer = SummaryWriter()
    path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', DataSetName)
    dataset = TUDataset(path, name=DataSetName)
    dataset = dataset.shuffle()
    n = len(dataset) // 10
    test_dataset = dataset[:n]
    train_dataset = dataset[n:]
    test_loader = DataLoader(test_dataset, batch_size=256)
    train_loader = DataLoader(train_dataset, batch_size=256)
    device = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu')
    model = Net(dataset).to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.0005)
    train_acc_record = []
    test_acc_record = []
    filename = DataSetName+'_'+str(times)+ ' GCN-4-node'

    for epoch in range(1, EPOCHES+1):
        loss = train(model,optimizer,train_dataset,train_loader,device)
        train_acc = tes(model,train_loader,device)
        train_acc_record.append(train_acc)
        test_acc = tes(model,test_loader,device)
        test_acc_record.append(test_acc)
        writer.add_scalars(filename+" Acc",{
            'trainAcc':train_acc,
            'testAcc':test_acc
        },epoch)
        writer.add_scalar(filename+' Loss/Train',loss,epoch)
        print('Epoch: {:03d}, Loss: {:.5f}, Train Acc: {:.5f}, Test Acc: {:.5f}'.
              format(epoch, loss, train_acc, test_acc))
    writer.flush()
    writer.close()

    train_acc_record = np.array(train_acc_record).reshape(-1,1)
    test_acc_record = np.array(test_acc_record).reshape(-1,1)
    record = np.hstack((train_acc_record,test_acc_record))
    np.savetxt(filename+'.txt',record)
    print(sum(test_acc_record)/EPOCHES)

if __name__ == '__main__':
    path = os.listdir("D:\\MasterStudents\\2020\ChenXin\\GDVshuffleTest\\data")
    for p in path:
        if os.path.isdir("D:\\MasterStudents\\2020\ChenXin\\GDVshuffleTest\\data\\" + p) and p.count('DD') != 0:
            print(p)
            PoolingWithGCN(p, 1, 800)
        if os.path.isdir("D:\\MasterStudents\\2020\ChenXin\\GDVshuffleTest\\data\\" + p) and p.count('MUTAG') != 0:
            print(p)
            PoolingWithGCN(p, 1, 3000)
        if os.path.isdir("D:\\MasterStudents\\2020\ChenXin\\GDVshuffleTest\\data\\" + p) and p.count('NCI1') != 0:
            print(p)
            PoolingWithGCN(p, 1, 1600)

